package com.hra.hrasittingposture.analysis;

import android.annotation.SuppressLint;
import android.content.Context;
import android.graphics.Bitmap;
import android.util.Log;

import androidx.annotation.NonNull;
import androidx.camera.core.ImageAnalysis;
import androidx.camera.core.ImageProxy;

import com.example.sittingposture.Recognition;
import com.example.sittingposture.Yolov5TFLiteDetector;
import com.hra.hrasittingposture.utils.ImageProcess;

import java.util.ArrayList;

public class FullImageAnalyse implements ImageAnalysis.Analyzer {
    private static final String TAG = "FullImageAnalyseDebug";
    private final ImageProcess imageProcess;
    private final Yolov5TFLiteDetector yolov5TFLiteDetector;
    private final Context context;
    private final Callback callback;
    private int counter = 0;

    public FullImageAnalyse(Context context, Yolov5TFLiteDetector yolov5TFLiteDetector, Callback callback) {
        this.context = context;
        this.yolov5TFLiteDetector = yolov5TFLiteDetector;
        this.callback = callback;
        this.imageProcess = new ImageProcess();
    }

    @Override
    public void analyze(@NonNull ImageProxy image) {
        counter++;
        if (counter == 120) {
            detect(image);
            counter = 0;
        }
        // 关闭 ImageProxy 以释放资源
        image.close();
    }

    @SuppressLint("DefaultLocale")
    private void detect(ImageProxy image) {
        long start = System.currentTimeMillis();
        // 获取 YUV 数据
        byte[][] yuvBytes = new byte[3][];
        ImageProxy.PlaneProxy[] planes = image.getPlanes();
        int imageHeight = image.getHeight();
        int imageWidth = image.getWidth();
        // 填充 YUV 数据
        imageProcess.fillBytes(planes, yuvBytes);
        int yRowStride = planes[0].getRowStride();
        int uvRowStride = planes[1].getRowStride();
        int uvPixelStride = planes[1].getPixelStride();
        // 创建 RGB 数组并转换 YUV 到 ARGB
        int[] rgbBytes = new int[imageHeight * imageWidth];
        imageProcess.YUV420ToARGB8888(
                yuvBytes[0],
                yuvBytes[1],
                yuvBytes[2],
                imageWidth,
                imageHeight,
                yRowStride,
                uvRowStride,
                uvPixelStride,
                rgbBytes);

        // 创建位图以供 YOLO 模型使用
        Bitmap imageBitmap = Bitmap.createBitmap(imageWidth, imageHeight, Bitmap.Config.ARGB_8888);
        imageBitmap.setPixels(rgbBytes, 0, imageWidth, 0, 0, imageWidth, imageHeight);

        // 进行 YOLO 识别
        ArrayList<Recognition> recognitions = yolov5TFLiteDetector.detect(imageBitmap);
        if (callback != null) {
            callback.getDetectResult(recognitions);
        }

        long end = System.currentTimeMillis();
        long costTime = (end - start);
        Log.d(TAG, "Processing time: " + costTime + " ms");
    }

    public interface Callback {
        void getDetectResult(ArrayList<Recognition> recognitions);
    }
}
